package com.cn.demo.util.threadpool;
import brave.Span;
import brave.Tracing;
import com.cn.demo.util.ApplicationContextUtil;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.*;

/**
 * 带有trace功能线程池
 */
@Slf4j
public class SdbThreadPoolUtil {

    private final int CORESIZE = Runtime.getRuntime().availableProcessors() * 2;

    private ThreadPoolExecutor threadPoolExecutor;

    @Getter
    private ExecutorService tracingExecutorService;

    private Tracing tracing;

    private String threadPoolGroupName;

    public SdbThreadPoolUtil(String threadPoolGroupName) {
        this.threadPoolGroupName = threadPoolGroupName;
        tracing = ApplicationContextUtil.getBean(Tracing.class) != null ?
                ApplicationContextUtil.getBean(Tracing.class) : Tracing.newBuilder().build();
        threadPoolExecutor = new LpThreadPoolExecutor(CORESIZE, CORESIZE * 2, 60, TimeUnit.SECONDS,
                new ArrayBlockingQueue<>(CORESIZE * 1000),new LpDefaultThreadFactory(threadPoolGroupName),
                new LpAbortPolicyWithReport(threadPoolGroupName));
        tracingExecutorService = tracing.currentTraceContext().executorService(threadPoolExecutor);

        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            tracingExecutorService.shutdown();
            log.info("SdbThreadPoolUtil {} shutdown success", threadPoolGroupName);
        }));
        log.info("SdbThreadPoolUtil init threadpool {} success", threadPoolGroupName);
    }


    public void execute(Runnable command) {
        printTraceInfo();
        Runnable wrapRunnable = tracing.currentTraceContext().wrap(command);
        tracingExecutorService.execute(wrapRunnable);
    }



    public <T> Future<T> submit(Callable<T> task) {
        printTraceInfo();
        Callable<T> wrapCallable = tracing.currentTraceContext().wrap(task);
        return tracingExecutorService.submit(wrapCallable);
    }

    private void printTraceInfo() {
        Span span = tracing.tracer().currentSpan() == null ? tracing.tracer().newTrace() : tracing.tracer().currentSpan();
        log.info("SdbThreadPoolUtil threadPoolGroupName={},execute traceid={}", threadPoolGroupName, span.context().traceId());
    }
}
